# stress_test.py

import sys
import os
import time
from datetime import datetime, timedelta
import random
from faker import Faker
import threading
from queue import Queue
from sqlalchemy import create_engine, Column, Integer, String, ForeignKey, CHAR, BIGINT
from sqlalchemy.orm import sessionmaker, Session, relationship, declarative_base
from sqlalchemy.schema import Table, MetaData, PrimaryKeyConstraint, ForeignKeyConstraint
from sqlalchemy.sql import text


# 1. ================== 配置 ==================
# 请修改为您自己的MySQL数据库连接信息
# 格式: "mysql+mysqlclient://<用户名>:<密码>@<主机>:<端口>/<数据库名>"
STRESS_TEST_DATABASE_URL = "mysql://root:123456@localhost:3306/hospital_stress_test_db"
BATCH_SIZE = 1000 # 每次提交的记录数
NUM_THREADS = 4 # 线程数

# 2. ================== 数据库设置 (与 app.py 保持一致的表结构定义) ==================
engine = create_engine(STRESS_TEST_DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

class Department(Base):
    __tablename__ = "departments"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(50), unique=True, nullable=False, index=True)
    doctors = relationship("Doctor", back_populates="department")

class Doctor(Base):
    __tablename__ = "doctors"
    id = Column(String(50), primary_key=True, index=True) # 工号
    name = Column(String(50), nullable=False)
    hashed_password = Column(String(100), nullable=False) # 压力测试简化为明文密码
    department_id = Column(Integer, ForeignKey("departments.id"), index=True)
    title = Column(String(50), nullable=True)
    department = relationship("Department", back_populates="doctors")
    __table_args__ = (
        ForeignKeyConstraint(
            ['department_id'], ['departments.id'], name='fk_doctors_department_id_stress'
        ),
    )

class Patient(Base):
    __tablename__ = "patients"
    id = Column(String(50), primary_key=True, index=True) # 病历号
    name = Column(String(50), nullable=False)
    gender = Column(CHAR(1), nullable=True)
    phone_number = Column(String(20), nullable=True, index=True)
    hashed_password = Column(String(100), nullable=False) # 压力测试简化为明文密码

class Appointment(Base):
    __tablename__ = "appointments"
    id = Column(Integer, primary_key=True, index=True)
    patient_id = Column(String(50), ForeignKey("patients.id"), index=True)
    doctor_id = Column(String(50), ForeignKey("doctors.id"), index=True)
    appointment_date = Column(String(20), default=lambda: datetime.now().strftime("%Y-%m-%d %H:%M"), index=True)
    status = Column(String(20), default="已预约")
    patient = relationship("Patient")
    doctor = relationship("Doctor")
    prescriptions = relationship("Prescription", back_populates="appointment")
    __table_args__ = (
        ForeignKeyConstraint(
            ['patient_id'], ['patients.id'], name='fk_appointments_patient_id_stress'
        ),
        ForeignKeyConstraint(
            ['doctor_id'], ['doctors.id'], name='fk_appointments_doctor_id_stress'
        ),
    )

class Prescription(Base):
    __tablename__ = "prescriptions"
    id = Column(Integer, primary_key=True, index=True)
    appointment_id = Column(Integer, ForeignKey("appointments.id"), index=True)
    item_name = Column(String(100))
    details = Column(String(255))
    appointment = relationship("Appointment", back_populates="prescriptions")
    __table_args__ = (
        ForeignKeyConstraint(
            ['appointment_id'], ['appointments.id'], name='fk_prescriptions_appointment_id_stress'
        ),
    )

# 3. ================== 数据生成与插入 ==================

def generate_departments(db: Session, num_departments=5):
    fake = Faker('zh_CN')
    departments = []
    for i in range(num_departments):
        dept = Department(name=fake.unique.word() + "科")
        departments.append(dept)
    db.add_all(departments)
    db.commit()
    return departments

def generate_doctors(db: Session, departments, num_doctors=100000):
    fake = Faker('zh_CN')
    doctors = []
    available_department_ids = [dept.id for dept in departments]
    for i in range(num_doctors):
        doctor_id = f"doc{i+1:05d}" # doc00001, doc00002...
        doc = Doctor(
            id=doctor_id,
            name=fake.name(),
            hashed_password=f"{i+1}123", # 简化密码
            department_id=random.choice(available_department_ids),
            title=random.choice(["主任医师", "副主任医师", "主治医师", "住院医师"])
        )
        doctors.append(doc)
        if len(doctors) % BATCH_SIZE == 0:
            db.add_all(doctors)
            db.commit()
            doctors = [] # 清空列表以便下批次
    if doctors: # 插入剩余的
        db.add_all(doctors)
        db.commit()
    return doctors

def generate_patients(db: Session, num_patients=1000000):
    fake = Faker('zh_CN')
    patients = []
    for i in range(num_patients):
        patient_id = f"pat{i+1:06d}" # pat000001, pat000002...
        pat = Patient(
            id=patient_id,
            name=fake.name(),
            gender=random.choice(['M', 'F']),
            phone_number=fake.phone_number(),
            hashed_password=f"{i+1}123" # 简化密码
        )
        patients.append(pat)
        if len(patients) % BATCH_SIZE == 0:
            db.add_all(patients)
            db.commit()
            patients = []
    if patients:
        db.add_all(patients)
        db.commit()
    return patients

def generate_appointments(db: Session, doctors, patients, num_appointments=5000000):
    fake = Faker('zh_CN')
    appointments = []
    all_doctor_ids = [doc.id for doc in doctors]
    all_patient_ids = [pat.id for pat in patients]
    for i in range(num_appointments):
        appointment_date = (datetime.now() + timedelta(days=random.randint(1, 30), hours=random.randint(9, 17), minutes=random.randint(0, 59))).strftime("%Y-%m-%d %H:%M")
        appt = Appointment(
            patient_id=random.choice(all_patient_ids),
            doctor_id=random.choice(all_doctor_ids),
            appointment_date=appointment_date,
            status=random.choice(["已预约", "已就诊", "已取消", "已错过"])
        )
        appointments.append(appt)
        if len(appointments) % BATCH_SIZE == 0:
            db.add_all(appointments)
            db.commit()
            appointments = []
            print(f"已生成并插入 {i+1}/{num_appointments} 条预约...")
    if appointments:
        db.add_all(appointments)
        db.commit()
    return appointments

def generate_prescriptions(db: Session, appointments, num_prescriptions=3000000):
    fake = Faker('zh_CN')
    prescriptions = []
    all_appointment_ids = [appt.id for appt in appointments]
    for i in range(num_prescriptions):
        pres = Prescription(
            appointment_id=random.choice(all_appointment_ids),
            item_name=fake.word() + "片",
            details=fake.sentence(nb_words=6)
        )
        prescriptions.append(pres)
        if len(prescriptions) % BATCH_SIZE == 0:
            db.add_all(prescriptions)
            db.commit()
            prescriptions = []
            print(f"已生成并插入 {i+1}/{num_prescriptions} 条处方...")
    if prescriptions:
        db.add_all(prescriptions)
        db.commit()
    return prescriptions

def generate_data(db: Session, num_departments=5, num_doctors=1000, num_patients=10000, num_appointments=50000, num_prescriptions=30000):
    fake = Faker('zh_CN') # 使用中文 Faker
    print("\n--- 生成部门数据 ---")
    departments = generate_departments(db, num_departments)
    print(f"生成并插入了 {len(departments)} 个部门。")

    print("\n--- 生成医生数据 ---")
    doctors = generate_doctors(db, departments, num_doctors)
    print(f"生成并插入了 {num_doctors} 名医生。")

    print("\n--- 生成病人数据 ---")
    patients = generate_patients(db, num_patients)
    print(f"生成并插入了 {num_patients} 名病人。")

    print("\n--- 生成预约数据 ---")
    appointments = generate_appointments(db, doctors, patients, num_appointments)
    print(f"生成并插入了 {num_appointments} 条预约。")

    print("\n--- 生成处方数据 ---")
    prescriptions = generate_prescriptions(db, appointments, num_prescriptions)
    print(f"生成并插入了 {num_prescriptions} 条处方。")

    print("\n--- 数据生成完毕 ---")


# 4. ================== 性能测试 ==================

def run_performance_tests(db: Session):
    print("\n--- 运行性能测试 ---")
    
    start_time = time.time()
    total_doctors = db.query(Doctor).count()
    end_time = time.time()
    print(f"查询总医生数量: {total_doctors} 条, 耗时: {end_time - start_time:.4f} 秒")

    start_time = time.time()
    total_patients = db.query(Patient).count()
    end_time = time.time()
    print(f"查询总病人数量: {total_patients} 条, 耗时: {end_time - start_time:.4f} 秒")
    
    start_time = time.time()
    total_appointments = db.query(Appointment).count()
    end_time = time.time()
    print(f"查询总预约数量: {total_appointments} 条, 耗时: {end_time - start_time:.4f} 秒")

    # 查询特定医生（例如第一个医生）的所有预约
    if db.query(Doctor).first():
        first_doctor_id = db.query(Doctor.id).first()[0]
        start_time = time.time()
        doctor_appointments = db.query(Appointment).filter(Appointment.doctor_id == first_doctor_id).all()
        end_time = time.time()
        print(f"查询医生 {first_doctor_id} 的 {len(doctor_appointments)} 条预约, 耗时: {end_time - start_time:.4f} 秒")
    else:
        print("无医生数据，跳过医生预约查询测试。")

    # 查询科室概况 (使用 app.py 中的 WITH CTE SQL)
    start_time = time.time()
    query_department_stats = text("""
        WITH DepartmentDoctorCount AS (
            SELECT
                d.department_id,
                dep.name AS department_name,
                COUNT(d.id) AS doctor_count
            FROM
                doctors AS d
            JOIN
                departments AS dep ON d.department_id = dep.id
            GROUP BY
                d.department_id, dep.name
        ),
        DepartmentActiveAppointmentCount AS (
            SELECT
                d.department_id,
                COUNT(a.id) AS active_appointment_count
            FROM
                appointments AS a
            JOIN
                doctors AS d ON a.doctor_id = d.id
            WHERE
                a.status = '已预约'
            GROUP BY
                d.department_id
        )
        SELECT
            ddc.department_name,
            ddc.doctor_count,
            COALESCE(daac.active_appointment_count, 0) AS total_active_appointments
        FROM
            DepartmentDoctorCount AS ddc
        LEFT JOIN
            DepartmentActiveAppointmentCount AS daac ON ddc.department_id = daac.department_id
        ORDER BY
            ddc.department_name;
    """)
    result_department_stats = db.execute(query_department_stats).fetchall()
    end_time = time.time()
    print(f"查询科室概况 ({len(result_department_stats)} 条记录), 耗时: {end_time - start_time:.4f} 秒")

    # 查询总体概览 (使用 app.py 中的聚集函数 SQL)
    start_time = time.time()
    query_overall_stats = text("""
        SELECT
            COUNT(DISTINCT d.id) AS total_doctors,
            COUNT(DISTINCT p.id) AS total_patients,
            COUNT(DISTINCT a.id) AS total_appointments,
            COUNT(DISTINCT pr.id) AS total_prescriptions,
            COALESCE(ROUND(SUM(CASE WHEN a.status = '已预约' THEN 1 ELSE 0 END) / COUNT(DISTINCT d.id), 2), 0) AS avg_active_appointments_per_doctor
        FROM
            doctors d,
            patients p,
            appointments a,
            prescriptions pr;
    """)
    result_overall_stats = db.execute(query_overall_stats).fetchone()
    end_time = time.time()
    print(f"查询总体概览, 耗时: {end_time - start_time:.4f} 秒")

    print("\n--- 性能测试完毕 ---")


# 5. ================== 主执行逻辑 ==================
if __name__ == "__main__":
    db = SessionLocal()
    try:
        print("--- 压力测试脚本启动 ---")

        # 尝试创建表 (如果数据库为空或表不存在)
        print("尝试在 hospital_stress_test_db 中创建表...")
        Base.metadata.create_all(bind=engine)
        print("表创建完成。")

        # 清空现有数据 (如果需要，以便每次测试都从干净状态开始)
        print("\n--- 清空现有数据（可选） ---")
        confirm = input("是否清空 hospital_stress_test_db 中所有现有数据？(y/n): ")
        if confirm.lower() == 'y':
            print("正在清空数据...")
            db.query(Prescription).delete()
            db.query(Appointment).delete()
            db.query(Doctor).delete()
            db.query(Patient).delete()
            db.query(Department).delete()
            db.commit()
            print("数据清空完毕。")
        else:
            print("跳过数据清空。如果数据库中有旧数据，测试结果可能不准确。")

        # 生成并插入数据
        start_gen_time = time.time()
        generate_data(db)
        end_gen_time = time.time()
        print(f"数据生成和插入总耗时: {end_gen_time - start_gen_time:.2f} 秒")

        # 运行性能测试
        run_performance_tests(db)

    except Exception as e:
        db.rollback()
        print(f"发生错误: {e}")
    finally:
        db.close()
        print("\n--- 压力测试脚本结束 ---") 